1806E - Tree Master - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
#define int long long
int T;
unordered_map<int,int>p[100005];
vector<int>ep[100005];
int dep[100005];
int n,q;
int a[100005];
int fa[100005];
int s[100005];
void dfs(int k,int sum){
	dep[k]=dep[fa[k]]+1;
	sum+=a[k]*a[k];
	s[k]=sum;
	for(int i=0;i<ep[k].size();i++){
		dfs(ep[k][i],sum);
	}
}
int f(int x,int y){
	if(x>y)swap(x,y);
	if(x==y)return s[x];
	if(dep[x]%3==0){
	
	int op=p[x][y]; 
	if(op)return op;
	return p[x][y]=f(fa[x],fa[y])+a[x]*a[y];
	}
	return f(fa[x],fa[y])+a[x]*a[y];
}
void solve(){
scanf("%lld %lld",&n,&q);
for(int i=1;i<=n;i++){
	scanf("%lld",&a[i]);
}
int t1,t2;
for(int i=2;i<=n;i++){
	scanf("%lld",&fa[i]);
	ep[fa[i]].push_back(i);
}
dfs(1,0);
for(int i=1;i<=q;i++){
	scanf("%lld %lld",&t1,&t2);
	printf("%lld\n",f(t1,t2));
}
}
signed main(){
T=1;
//ios::sync_with_stdio(0);
//cin.tie(0);
//cin>>T;
while(T--){
solve();
 
}
 
 
 
 
return 0;
}